// Copyright 2022 jeff.li. and/or its affiliates.


/*!
 * \file src/printer/doc.cc
 * \brief Doc ADT used for pretty printing.
 *
 *  Reference: Philip Wadler. A Prettier Printer. Journal of Functional Programming'98
 */
#include "doc.h"

#include <sstream>
#include <vector>

#include <tbir/runtime/container/unicode.h>
#include <tbir/runtime/memory.h>
#include <tbir/runtime/str_escape.h>
#include <tbir/runtime/utf8_util.h>

namespace tbir::printer {

    /*!
     * \brief Represent a piece of text in the doc.
     */
    class DocTextNode : public DocAtomNode {
    public:
        /*! \brief The str content in the text. */
        runtime::String str;

        explicit DocTextNode(runtime::String str_val) : str(str_val) {
        }

        static constexpr const char *_type_key = "printer.DocText";
        TBIR_DECLARE_FINAL_OBJECT_INFO(DocTextNode, DocAtomNode);
    };

    TBIR_REGISTER_OBJECT_TYPE(DocTextNode);

    class DocText : public DocAtom {
    public:
        explicit DocText(runtime::String str) {
            if (str.find_first_of("\t\n") != str.npos) {
                MXLOG(WARNING) << "text node: '" << str << "' should not has tab or newline.";
            }
            data_ = runtime::make_object<DocTextNode>(str);
        }

        TBIR_DEFINE_OBJECT_REF_METHODS(DocText, DocAtom, DocTextNode);
    };

    /*!
     * \brief Represent a line breaker in the doc.
     */
    class DocLineNode : public DocAtomNode {
    public:
        /*! \brief The amount of indent in newline. */
        int indent;

        explicit DocLineNode(int indent) : indent(indent) {
        }

        static constexpr const char *_type_key = "printer.DocLine";
        TBIR_DECLARE_FINAL_OBJECT_INFO(DocLineNode, DocAtomNode);
    };

    TBIR_REGISTER_OBJECT_TYPE(DocLineNode);

    class DocLine : public DocAtom {
    public:
        explicit DocLine(int indent) {
            data_ = runtime::make_object<DocLineNode>(indent);
        }

        TBIR_DEFINE_OBJECT_REF_METHODS(DocLine, DocAtom, DocLineNode);
    };

    // DSL function implementations
    Doc &Doc::operator<<(const Doc &right) {
        MXCHECK(this != &right);
        this->stream_.insert(this->stream_.end(), right.stream_.begin(), right.stream_.end());
        return *this;
    }

    Doc &Doc::operator<<(runtime::String right) {
        return *this << DocText(right);
    }

    Doc &Doc::operator<<(const DocAtom &right) {
        this->stream_.push_back(right);
        return *this;
    }

    runtime::String Doc::str() {
        std::ostringstream os;
        for (auto atom : this->stream_) {
            if (auto *text = atom.as<DocTextNode>()) {
                os << text->str;
            } else if (auto *line = atom.as<DocLineNode>()) {
                os << "\n" << runtime::String(line->indent, ' ');
            } else {
                MXLOG(FATAL) << "do not expect type " << atom->GetTypeKey();
            }
        }
        return os.str();
    }

    Doc Doc::NewLine(int indent) {
        return Doc() << DocLine(indent);
    }

    Doc Doc::Text(runtime::String text) {
        return Doc() << DocText(text);
    }

    Doc Doc::RawText(runtime::String text) {
        return Doc() << DocAtom(runtime::make_object<DocTextNode>(text));
    }

    Doc Doc::Indent(int indent, Doc doc) {
        for (size_t i = 0; i < doc.stream_.size(); ++i) {
            if (auto *line = doc.stream_[i].as<DocLineNode>()) {
                doc.stream_[i] = DocLine(indent + line->indent);
            }
        }
        return doc;
    }

    Doc Doc::StrLiteral(const runtime::String &value, const runtime::String &quote) {
        Doc doc;
        return doc << quote << runtime::BytesEscape(value) << quote;
    }

    Doc Doc::PyBoolLiteral(bool value) {
        if (value) {
            return Doc::Text("True");
        } else {
            return Doc::Text("False");
        }
    }

    Doc Doc::Brace(runtime::String open, const Doc &body, runtime::String close, int indent) {
        Doc doc;
        doc << open;
        doc << Indent(indent, NewLine() << body) << NewLine();
        doc << close;
        return doc;
    }

    Doc Doc::Concat(const std::vector<Doc> &vec, const Doc &sep) {
        Doc seq;
        if (vec.size() != 0) {
            if (vec.size() == 1)
                return vec[0];
            seq << vec[0];
            for (size_t i = 1; i < vec.size(); ++i) {
                seq << sep << vec[i];
            }
        }
        return seq;
    }

}  // namespace tbir::printer
